#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2025/2/5 16:19
# @Author  : Fancy
# @FileName: postgres.py
# @Software: PyCharm
# @runstate    ：
# clients/postgres.py
import psycopg2
from db_compare.clients.base import BaseClient

class PGClient(BaseClient):
    def __init__(self, config):
        self.config = config
        print(f"PGClient 初始化, config: {self.config}")
    """PostgreSQL 客户端实现"""
    def get_conn(self):
        return psycopg2.connect(
            host=self.config.connect_args["host"],
            user=self.config.connect_args["user"],
            password=self.config.connect_args["password"],
            database=self.config.connect_args["database"],
            port=int (self.config.connect_args["port"])
        )

    def fetch_stream (self, table: str, batch_size: int = 1000):
        """流式获取数据"""
        offset = 0
        conn = self.get_conn ()
        try:
            while True:
                with conn.cursor () as cursor:
                    cursor.execute (f"SELECT * FROM {table} LIMIT {batch_size} OFFSET {offset}")
                    batch = cursor.fetchall ()
                    if not batch:
                        break
                    yield batch
                    offset += batch_size
                    print (f"获取表 {table} 的数据流，每批 {batch_size} 条")
        finally:
            conn.close ()